import torch
import torch.nn as nn


class Predictor(nn.Module):
    def __init__(self, 
                 meteor_hidden_size,
                 pollutant_hidden_size,
                 num_future_steps=6,
                 num_sites=39,
                 d_model=128,
                 time_shift_emb_dim=16,
                 nhead=8,
                 dropout=0.1):
        super().__init__()
        self.num_future_steps = num_future_steps
        self.num_sites = num_sites
        self.d_model = d_model

        self.meteor_proj = nn.Linear(meteor_hidden_size, d_model)
        self.time_shift_embedding = nn.Embedding(10, time_shift_emb_dim)
        self.shift_to_q = nn.Linear(time_shift_emb_dim, d_model)
        self.shift_to_kv = nn.Linear(time_shift_emb_dim, d_model)
        self.future_time_queries = nn.Parameter(torch.randn(num_future_steps, d_model))
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=nhead,
            dropout=dropout,
            batch_first=True
        )
        self.meteor_out_proj = nn.Linear(d_model, pollutant_hidden_size)

        self.regressor = nn.Sequential(
            nn.Linear(pollutant_hidden_size, 64),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(64, 1)
        )
        self.norm = nn.LayerNorm(d_model)


    def forward(self, 
                meteor_hidden,
                pollutant_hidden,
                time_shift): 

        B = pollutant_hidden.size(0)
        
        meteor_proj = self.meteor_proj(meteor_hidden) 
        meteor_proj = self.norm(meteor_proj) 

        ts_emb = self.time_shift_embedding(time_shift)
        ts_q_bias = self.shift_to_q(ts_emb) 
        ts_kv_bias = self.shift_to_kv(ts_emb) 

        queries = self.future_time_queries.unsqueeze(0).expand(B, -1, -1)
        queries = queries + ts_q_bias.unsqueeze(1)

        kv = meteor_proj + ts_kv_bias.unsqueeze(1) 

        attn_out, attn_weights = self.cross_attn(
            query=queries, 
            key=kv,
            value=kv,
            attn_mask=None,
            key_padding_mask=None
        ) 

        meteor_context = self.meteor_out_proj(attn_out)

        meteor_broadcast = meteor_context.unsqueeze(2).expand(-1, -1, self.num_sites, -1) 

        fused = pollutant_hidden + meteor_broadcast 

        output = self.regressor(fused)
        output = output.squeeze(-1)

        return output, attn_weights